#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @FileName  :q1_test3.py
# @Time      :2024/9/5 23:47
# @Author    :YKW
from q1_test2 import *
import sys
import pandas as pd

# TODO 目前还没特殊处理 龙头较长的问题
# TODO 现在数据还是不对  --解决

pos = [[] for _ in range(0, 301)]
vel = [[1] for _ in range(0, 301)]
pos[0].append((880, 0))
# 先把龙头的位置跑出来
for time in range(1, 301):
    pos[time].append(tuple(calculate_new_coordinate(pos[0][0], 1, time)))

for time in range(1, 301):
    # 进度条
    percentage = round(time / 301 * 100)
    print("\r进度: {}%: ".format(percentage), "▓" * (percentage // 2), end="")
    sys.stdout.flush()

    # 算位置、速度
    for point in range(1, 234):
        # 根据前一个点，此时刻的位置/速度，算出下一个点位置/速度
        x_b, y_b, v_b = calculate_pos(pos[time][-1][0], pos[time][-1][1], vel[time][-1],
                                      (lambda pit: 165 if pit == 1 else 286)(point))
        pos[time].append((x_b, y_b))
        vel[time].append(v_b)

# 创建一个 ExcelWriter 对象
with pd.ExcelWriter('result1.xlsx') as writer:
    # 处理 pos 数据并写入 "位置" 工作簿
    pos_data = {
        '时间': [f'{i}s' for i in range(301)]
    }

    # 自动判断 pos 的大小
    max_points = max(len(p) for p in pos)
    for point in range(max_points):
        pos_data[f'第{point + 1}点x'] = [pos[i][point][0] if len(pos[i]) > point and len(pos[i][point]) > 0 else None
                                         for i in range(301)]
        pos_data[f'第{point + 1}点y'] = [pos[i][point][1] if len(pos[i]) > point and len(pos[i][point]) > 1 else None
                                         for i in range(301)]

    pos_df = pd.DataFrame(pos_data)
    pos_df.to_excel(writer, sheet_name='位置', index=False)

    # 处理 vel 数据并写入 "速度" 工作簿
    vel_data = {
        '时间': [f'{i}s' for i in range(301)]
    }

    # 自动判断 vel 的大小
    max_points = max(len(v) for v in vel)
    for point in range(max_points):
        vel_data[f'第{point + 1}点(m/s)'] = [vel[i][point] if len(vel[i]) > point else None for i in range(301)]

    vel_df = pd.DataFrame(vel_data)
    # vel_df=vel_df.transpose()
    vel_df.to_excel(writer, sheet_name='速度', index=False)
